package io.openbas.rest.inject.service;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.openbas.database.model.*;
import io.openbas.database.repository.AgentRepository;
import io.openbas.database.repository.InjectExpectationRepository;
import io.openbas.database.repository.InjectRepository;
import io.openbas.rest.exception.ElementNotFoundException;
import io.openbas.rest.finding.FindingService;
import io.openbas.rest.inject.form.InjectExecutionAction;
import io.openbas.rest.inject.form.InjectExecutionInput;
import io.openbas.rest.inject.form.InjectExpectationUpdateInput;
import io.openbas.service.InjectExpectationService;
import jakarta.annotation.Resource;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.dao.DataIntegrityViolationException;
import org.springframework.stereotype.Service;

@RequiredArgsConstructor
@Service
@Slf4j
public class InjectExecutionService {

  private final InjectRepository injectRepository;
  private final InjectExpectationRepository injectExpectationRepository;
  private final InjectExpectationService injectExpectationService;
  private final AgentRepository agentRepository;
  private final InjectStatusService injectStatusService;
  private final FindingService findingService;
  private final StructuredOutputUtils structuredOutputUtils;

  @Resource protected ObjectMapper mapper;

  public void handleInjectExecutionCallback(
      String injectId, String agentId, InjectExecutionInput input) {
    Inject inject = null;

    try {
      inject = loadInjectOrThrow(injectId);
      // issue/3550: added this condition to ensure we only update statuses if the inject is in a
      // coherent state.
      // This prevents issues where the PENDING status took more time to persist than it took for
      // the agent to send the complete action.
      // FIXME: At the moment, this whole function is only called by our implant. These implant are
      // launched with the async value to true, which force the implant to go from EXECUTING to
      // PENDING, before going to EXECUTED.
      // So if in the future, this function is called to update a synchronous inject, we will need
      // to find a way to get the async boolean somehow and add it to this condition.
      if (input.getAction().equals(InjectExecutionAction.complete)
          && (inject.getStatus().isEmpty()
              || !inject.getStatus().get().getName().equals(ExecutionStatus.PENDING))) {
        // If we receive a status update with a terminal state status, we must first check that the
        // current status is in the PENDING state
        log.warn(
            String.format(
                "Received a complete action for inject %s with status %s, but current status is not PENDING",
                injectId, inject.getStatus().map(is -> is.getName().toString()).orElse("unknown")));
        throw new DataIntegrityViolationException(
            "Cannot complete inject that is not in PENDING state");
      }
      Agent agent = loadAgentIfPresent(agentId);

      Set<OutputParser> outputParsers = structuredOutputUtils.extractOutputParsers(inject);
      Optional<ObjectNode> structuredOutput =
          structuredOutputUtils.computeStructuredOutput(outputParsers, input);

      processInjectExecution(inject, agent, input, outputParsers, structuredOutput);
    } catch (ElementNotFoundException | JsonProcessingException e) {
      handleInjectExecutionError(inject, e);
    }
  }

  /** Processes the execution of an inject by updating its status and extracting findings. */
  private void processInjectExecution(
      Inject inject,
      Agent agent,
      InjectExecutionInput input,
      Set<OutputParser> outputParsers,
      Optional<ObjectNode> structuredOutput) {

    ObjectNode structured = structuredOutput.orElse(null);
    injectStatusService.updateInjectStatus(agent, inject, input, structured);

    if (structured == null) {
      return;
    }

    if (agent != null) {
      // validate vulnerability expectations
      checkCveExpectation(outputParsers, structured, inject, agent);

      // Extract findings from structured outputs generated by the output parsers specified in the
      // payload, typically derived from the raw output of the implant execution.
      findingService.extractFindingsFromOutputParsers(inject, agent, outputParsers, structured);
    } else {
      // Structured output directly provided (e.g., from injectors)
      findingService.extractFindingsFromInjectorContract(inject, structured);
    }
  }

  /**
   * Checks output parsers of an agent and updates the scores of vulnerability expectations
   * accordingly
   *
   * @param outputParsers
   * @param structuredOutput
   * @param inject
   * @param agent
   */
  public void checkCveExpectation(
      Set<OutputParser> outputParsers, ObjectNode structuredOutput, Inject inject, Agent agent) {
    List<InjectExpectation> injectExpectations = new ArrayList<>();

    inject.getExpectations().stream()
        .filter(injectExpectation -> injectExpectation.getAgent() != null)
        .filter(injectExpectation -> injectExpectation.getAgent().getId().equals(agent.getId()))
        .forEach(
            expectation -> {
              if (expectation.getType() == InjectExpectation.EXPECTATION_TYPE.VULNERABILITY) {
                injectExpectations.add(expectation);
              }
            });

    if (!injectExpectations.isEmpty()) {
      InjectExpectationResult injectExpectationResult =
          InjectExpectationResult.builder()
              .sourceId("acab8214-0379-448a-a575-05e9d934eadd")
              .date(String.valueOf(Instant.now()))
              .sourceType("openbas_expectations_vulnerability_manager")
              .sourceName("Expectations Vulnerability Manager")
              .score(0.0)
              .result("Vulnerable")
              .metadata(null)
              .build();
      outputParsers.forEach(
          outputParser -> {
            outputParser
                .getContractOutputElements()
                .forEach(
                    contractOutputElement -> {
                      if (contractOutputElement.getType().equals(ContractOutputType.CVE)) {
                        JsonNode jsonNode = structuredOutput.get(contractOutputElement.getKey());
                        if (jsonNode != null) {
                          if (!jsonNode.isEmpty()) {
                            injectExpectations.forEach(
                                expectation -> {
                                  expectation.setScore(0.0);
                                  expectation.setResults(List.of(injectExpectationResult));
                                });
                          } else {
                            injectExpectations.forEach(
                                expectation -> {
                                  expectation.setScore(expectation.getExpectedScore());
                                  injectExpectationResult.setResult("Not vulnerable");
                                  injectExpectationResult.setScore(expectation.getExpectedScore());
                                  expectation.setResults(List.of(injectExpectationResult));
                                });
                          }
                          injectExpectationRepository.saveAll(injectExpectations);
                          validateResultForAsset(injectExpectations, injectExpectationResult);
                        }
                      }
                    });
          });
    }
  }

  public void validateResultForAsset(
      List<InjectExpectation> injectExpectations, InjectExpectationResult injectExpectationResult) {
    injectExpectations.forEach(
        injectExpectation -> {
          injectExpectationService.updateInjectExpectation(
              injectExpectation.getId(),
              InjectExpectationUpdateInput.builder()
                  .collectorId(injectExpectationResult.getSourceId())
                  .result(injectExpectationResult.getResult())
                  .isSuccess(injectExpectationResult.getScore() != 0.0)
                  .build());
        });
  }

  private Agent loadAgentIfPresent(String agentId) {
    return (agentId == null)
        ? null
        : agentRepository
            .findById(agentId)
            .orElseThrow(() -> new ElementNotFoundException("Agent not found: " + agentId));
  }

  private Inject loadInjectOrThrow(String injectId) {
    return injectRepository
        .findById(injectId)
        .orElseThrow(() -> new ElementNotFoundException("Inject not found: " + injectId));
  }

  private void handleInjectExecutionError(Inject inject, Exception e) {
    log.error(e.getMessage(), e);
    if (inject != null) {
      inject
          .getStatus()
          .ifPresent(
              status -> {
                ExecutionTrace trace =
                    new ExecutionTrace(
                        status,
                        ExecutionTraceStatus.ERROR,
                        null,
                        e.getMessage(),
                        ExecutionTraceAction.COMPLETE,
                        null,
                        Instant.now());
                status.addTrace(trace);
              });
      injectRepository.save(inject);
    }
  }
}
